msg_tool\scripts\kirikiri/
scn.rs

1//! Kirikiri Scene File (.scn)
2use super::mdf::Mdf;
3use crate::ext::io::*;
4use crate::ext::mutex::*;
5use crate::ext::psb::*;
6use crate::scripts::base::*;
7use crate::types::*;
8use crate::utils::encoding::*;
9use anyhow::Result;
10use emote_psb::PsbReader;
11use fancy_regex::Regex;
12use std::collections::{HashMap, HashSet};
13use std::io::{Read, Seek};
14use std::path::Path;
15use std::sync::{Arc, Mutex};
16
17#[derive(Debug)]
18/// Kirikiri Scene Script Builder
19pub struct ScnScriptBuilder {}
20
21impl ScnScriptBuilder {
22    /// Creates a new instance of `ScnScriptBuilder`
23    pub fn new() -> Self {
24        Self {}
25    }
26}
27
28impl ScriptBuilder for ScnScriptBuilder {
29    fn default_encoding(&self) -> Encoding {
30        Encoding::Utf8
31    }
32
33    fn build_script(
34        &self,
35        buf: Vec<u8>,
36        filename: &str,
37        _encoding: Encoding,
38        _archive_encoding: Encoding,
39        config: &ExtraConfig,
40        _archive: Option<&Box<dyn Script>>,
41    ) -> Result<Box<dyn Script>> {
42        Ok(Box::new(ScnScript::new(
43            MemReader::new(buf),
44            filename,
45            config,
46        )?))
47    }
48
49    fn build_script_from_file(
50        &self,
51        filename: &str,
52        _encoding: Encoding,
53        _archive_encoding: Encoding,
54        config: &ExtraConfig,
55        _archive: Option<&Box<dyn Script>>,
56    ) -> Result<Box<dyn Script>> {
57        if filename == "-" {
58            let data = crate::utils::files::read_file(filename)?;
59            Ok(Box::new(ScnScript::new(
60                MemReader::new(data),
61                filename,
62                config,
63            )?))
64        } else {
65            let f = std::fs::File::open(filename)?;
66            let reader = std::io::BufReader::new(f);
67            Ok(Box::new(ScnScript::new(reader, filename, config)?))
68        }
69    }
70
71    fn build_script_from_reader(
72        &self,
73        reader: Box<dyn ReadSeek>,
74        filename: &str,
75        _encoding: Encoding,
76        _archive_encoding: Encoding,
77        config: &ExtraConfig,
78        _archive: Option<&Box<dyn Script>>,
79    ) -> Result<Box<dyn Script>> {
80        Ok(Box::new(ScnScript::new(reader, filename, config)?))
81    }
82
83    fn extensions(&self) -> &'static [&'static str] {
84        &["scn"]
85    }
86
87    fn script_type(&self) -> &'static ScriptType {
88        &ScriptType::KirikiriScn
89    }
90
91    fn is_this_format(&self, filename: &str, buf: &[u8], buf_len: usize) -> Option<u8> {
92        if Path::new(filename)
93            .file_name()
94            .map(|name| {
95                name.to_ascii_lowercase()
96                    .to_string_lossy()
97                    .ends_with(".scn")
98            })
99            .unwrap_or(false)
100            && buf_len >= 4
101            && buf.starts_with(b"PSB\0")
102        {
103            return Some(255);
104        }
105        None
106    }
107}
108
109#[derive(Debug)]
110/// Kirikiri Scene Script
111pub struct ScnScript {
112    psb: VirtualPsbFixed,
113    language_index: usize,
114    languages: Option<Arc<Vec<String>>>,
115    export_chat: bool,
116    filename: String,
117    chat_key: Option<Vec<String>>,
118    chat_json: Option<Arc<HashMap<String, HashMap<String, (String, usize)>>>>,
119    custom_yaml: bool,
120    title: bool,
121    chat_multilang: bool,
122    insert_language: bool,
123}
124
125impl ScnScript {
126    /// Creates a new `ScnScript` from the given reader and filename
127    ///
128    /// * `reader` - The reader containing the PSB or MDF data
129    /// * `filename` - The name of the file (used for error reporting and extension detection)
130    /// * `config` - Extra configuration options
131    pub fn new<R: Read + Seek>(
132        mut reader: R,
133        filename: &str,
134        config: &ExtraConfig,
135    ) -> Result<Self> {
136        let mut header = [0u8; 4];
137        reader.read_exact(&mut header)?;
138        if &header == b"mdf\0" {
139            let mut data = Vec::new();
140            reader.read_to_end(&mut data)?;
141            let decoded = Mdf::unpack(MemReaderRef::new(&data))?;
142            return Self::new(MemReader::new(decoded), filename, config);
143        }
144        reader.rewind()?;
145        let psb = PsbReader::open_psb_v2(reader)?;
146        Ok(Self {
147            psb: psb.to_psb_fixed(),
148            language_index: config.kirikiri_language_index.unwrap_or(0),
149            languages: config.kirikiri_languages.clone(),
150            export_chat: config.kirikiri_export_chat,
151            filename: filename.to_string(),
152            chat_key: config.kirikiri_chat_key.clone(),
153            chat_json: config.kirikiri_chat_json.clone(),
154            custom_yaml: config.custom_yaml,
155            title: config.kirikiri_title,
156            chat_multilang: config.kirikiri_chat_multilang,
157            insert_language: config.kirikiri_language_insert,
158        })
159    }
160}
161
162impl Script for ScnScript {
163    fn default_output_script_type(&self) -> OutputScriptType {
164        OutputScriptType::Json
165    }
166
167    fn default_format_type(&self) -> FormatOptions {
168        FormatOptions::None
169    }
170
171    fn is_output_supported(&self, _: OutputScriptType) -> bool {
172        true
173    }
174
175    fn custom_output_extension<'a>(&'a self) -> &'a str {
176        if self.custom_yaml { "yaml" } else { "json" }
177    }
178
179    fn extract_messages(&self) -> Result<Vec<Message>> {
180        let mut messages = Vec::new();
181        let root = self.psb.root();
182        let scenes = root
183            .get_value("scenes")
184            .ok_or(anyhow::anyhow!("scenes not found"))?;
185        let scenes = match scenes {
186            PsbValueFixed::List(list) => list,
187            _ => return Err(anyhow::anyhow!("scenes is not a list")),
188        };
189        let language = if self.language_index != 0 {
190            let index = self.language_index - 1;
191            if let Some(lang) = root["languages"][index].as_str() {
192                Some(lang.to_owned())
193            } else if let Some(languages) = self.languages.as_ref() {
194                if index < languages.len() {
195                    eprintln!(
196                        "WARN: Language code not found in PSB, using from config. Chat messages may not be extracted correctly."
197                    );
198                    crate::COUNTER.inc_warning();
199                    Some(languages[index].to_owned())
200                } else {
201                    None
202                }
203            } else {
204                None
205            }
206        } else {
207            None
208        };
209        if self.language_index != 0 && language.is_none() {
210            eprintln!(
211                "WARN: Language index is set but language code not found in PSB. Chat messages may not be extracted correctly."
212            );
213            crate::COUNTER.inc_warning();
214        }
215        let mut comu = if self.export_chat {
216            Some(ExportMes::new(
217                self.chat_key
218                    .clone()
219                    .unwrap_or(vec!["comumode".to_string()]),
220                if self.chat_multilang {
221                    language.clone()
222                } else {
223                    None
224                },
225            ))
226        } else {
227            None
228        };
229        for (i, oscene) in scenes.iter().enumerate() {
230            let scene = match oscene {
231                PsbValueFixed::Object(obj) => obj,
232                _ => return Err(anyhow::anyhow!("scene at index {} is not an object", i)),
233            };
234            if self.title {
235                if let Some(title) = scene["title"].as_str() {
236                    messages.push(Message {
237                        name: None,
238                        message: title.to_string(),
239                    });
240                }
241                if scene["title"].is_list() {
242                    if let Some(title) = scene["title"][self.language_index].as_str() {
243                        messages.push(Message {
244                            name: None,
245                            message: title.to_string(),
246                        });
247                    }
248                }
249            }
250            if let Some(PsbValueFixed::List(texts)) = scene.get_value("texts") {
251                for (j, text) in texts.iter().enumerate() {
252                    if let PsbValueFixed::List(text) = text {
253                        let values = text.values();
254                        if values.len() <= 1 {
255                            continue; // Skip if there are not enough values
256                        }
257                        let name = &values[0];
258                        let name = match name {
259                            PsbValueFixed::String(s) => Some(s),
260                            PsbValueFixed::Null => None,
261                            _ => return Err(anyhow::anyhow!("name is not a string or null")),
262                        };
263                        let mut display_name;
264                        let mut message;
265                        if matches!(values[1], PsbValueFixed::List(_)) {
266                            display_name = None;
267                            message = &values[1];
268                        } else {
269                            if values.len() <= 2 {
270                                continue; // Skip if there is no message
271                            }
272                            display_name = match &values[1] {
273                                PsbValueFixed::String(s) => Some(s),
274                                PsbValueFixed::Null => None,
275                                _ => {
276                                    return Err(anyhow::anyhow!(
277                                        "display name is not a string or null at {i},{j}"
278                                    ));
279                                }
280                            };
281                            message = &values[2];
282                        }
283                        if matches!(message, PsbValueFixed::List(_)) {
284                            let tmp = message;
285                            if let PsbValueFixed::List(list) = tmp {
286                                if list.len() > self.language_index {
287                                    if let PsbValueFixed::List(data) =
288                                        &list.values()[self.language_index]
289                                    {
290                                        if data.len() >= 2 {
291                                            let data = data.values();
292                                            display_name = match &data[0] {
293                                                PsbValueFixed::String(s) => Some(s),
294                                                PsbValueFixed::Null => None,
295                                                _ => {
296                                                    return Err(anyhow::anyhow!(
297                                                        "display name is not a string or null at {i},{j}"
298                                                    ));
299                                                }
300                                            };
301                                            message = &data[1];
302                                        }
303                                    }
304                                }
305                            }
306                        }
307                        if let PsbValueFixed::String(message) = message {
308                            match name {
309                                Some(name) => {
310                                    let name = match display_name {
311                                        Some(name) => name.string(),
312                                        None => name.string(),
313                                    };
314                                    let message = message.string();
315                                    messages.push(Message {
316                                        name: Some(name.to_string()),
317                                        message: message.replace("\\n", "\n"),
318                                    });
319                                }
320                                None => {
321                                    let message = message.string();
322                                    messages.push(Message {
323                                        name: None,
324                                        message: message.replace("\\n", "\n"),
325                                    });
326                                }
327                            }
328                        }
329                    }
330                }
331            }
332            if let Some(PsbValueFixed::List(selects)) = scene.get_value("selects") {
333                for select in selects.iter() {
334                    if let PsbValueFixed::Object(select) = select {
335                        let mut text = None;
336                        if let Some(PsbValueFixed::List(language)) = select.get_value("language") {
337                            if language.len() > self.language_index {
338                                let v = &language.values()[self.language_index];
339                                if let PsbValueFixed::Object(v) = v {
340                                    text = match v.get_value("text") {
341                                        Some(PsbValueFixed::String(s)) => Some(s),
342                                        Some(PsbValueFixed::Null) => None,
343                                        None => None,
344                                        _ => {
345                                            return Err(anyhow::anyhow!(
346                                                "select text is not a string or null"
347                                            ));
348                                        }
349                                    }
350                                }
351                            }
352                        }
353                        if text.is_none() {
354                            text = match select.get_value("text") {
355                                Some(PsbValueFixed::String(s)) => Some(s),
356                                Some(PsbValueFixed::Null) => None,
357                                None => None,
358                                _ => {
359                                    return Err(anyhow::anyhow!(
360                                        "select text is not a string or null"
361                                    ));
362                                }
363                            };
364                        }
365                        if let Some(text) = text {
366                            let text = text.string();
367                            messages.push(Message {
368                                name: None,
369                                message: text.replace("\\n", "\n"),
370                            });
371                        }
372                    }
373                }
374            }
375            comu.as_mut().map(|c| c.export(&oscene));
376        }
377        if let Some(comu) = comu {
378            if !comu.messages.is_empty() {
379                let mut pb = std::path::PathBuf::from(&self.filename);
380                let key = self
381                    .chat_key
382                    .clone()
383                    .unwrap_or(vec!["comumode".to_string()])
384                    .join("_");
385                let filename = pb
386                    .file_stem()
387                    .map(|s| s.to_string_lossy())
388                    .unwrap_or(std::borrow::Cow::from(&key));
389                pb.set_file_name(format!("{}_{}.json", filename, key));
390                match std::fs::File::create(&pb) {
391                    Ok(mut f) => {
392                        let messages: Vec<String> = comu.messages.into_iter().collect();
393                        if let Err(e) = serde_json::to_writer_pretty(&mut f, &messages) {
394                            eprintln!("Failed to write chat messages to {}: {:?}", pb.display(), e);
395                            crate::COUNTER.inc_warning();
396                        }
397                    }
398                    Err(e) => {
399                        eprintln!(
400                            "Failed to create chat messages file {}: {:?}",
401                            pb.display(),
402                            e
403                        );
404                        crate::COUNTER.inc_warning();
405                    }
406                }
407            }
408        }
409        Ok(messages)
410    }
411
412    fn import_messages<'a>(
413        &'a self,
414        messages: Vec<Message>,
415        file: Box<dyn WriteSeek + 'a>,
416        _filename: &str,
417        _encoding: Encoding,
418        replacement: Option<&'a ReplacementTable>,
419    ) -> Result<()> {
420        let mut mes = messages.iter();
421        let mut cur_mes = mes.next();
422        let mut psb = self.psb.clone();
423        let root = psb.root_mut();
424        if let Some(lang) = &self.languages {
425            let lang = (**lang).clone();
426            root["languages"] = PsbValueFixed::List(PsbListFixed {
427                values: lang
428                    .into_iter()
429                    .map(|s| PsbValueFixed::String(s.into()))
430                    .collect(),
431            });
432        }
433        let language = if self.language_index != 0 {
434            let index = self.language_index - 1;
435            if let Some(lang) = root["languages"][index].as_str() {
436                Some(lang.to_owned())
437            } else {
438                eprintln!(
439                    "WARN: language code not found in PSB. Some functions may not work correctly. Use --kirikiri-languages to specify language codes."
440                );
441                crate::COUNTER.inc_warning();
442                None
443            }
444        } else {
445            None
446        };
447        let ori_lang = if self.insert_language && self.language_index == 0 {
448            if let Some(lang) = root["languages"][0].as_str() {
449                Some(lang.to_owned())
450            } else {
451                None
452            }
453        } else {
454            None
455        };
456        let scenes = &mut root["scenes"];
457        if !scenes.is_list() {
458            return Err(anyhow::anyhow!("scenes is not an array"));
459        }
460        let comu = self.chat_json.as_ref().map(|json| {
461            ImportMes::new(
462                json,
463                replacement,
464                self.chat_key
465                    .clone()
466                    .unwrap_or(vec!["comumode".to_string()]),
467                if self.chat_multilang {
468                    language.clone()
469                } else {
470                    None
471                },
472                self.filename.clone(),
473                if self.chat_multilang {
474                    ori_lang.clone()
475                } else {
476                    None
477                },
478            )
479        });
480        for (i, scene) in scenes.members_mut().enumerate() {
481            if !scene.is_object() {
482                return Err(anyhow::anyhow!("scene at {} is not an object", i));
483            }
484            if self.title {
485                if scene["title"].is_string() {
486                    let m = match cur_mes {
487                        Some(m) => m,
488                        None => {
489                            return Err(anyhow::anyhow!(
490                                "No enough messages. (title at scene {i})"
491                            ));
492                        }
493                    };
494                    let mut title = m.message.clone();
495                    if let Some(replacement) = replacement {
496                        for (key, value) in replacement.map.iter() {
497                            title = title.replace(key, value);
498                        }
499                    }
500                    if self.language_index == 0 {
501                        if self.insert_language {
502                            let ori_title = scene["title"].as_str().unwrap_or("").to_string();
503                            scene["title"].push_member(title);
504                            scene["title"].push_member(ori_title);
505                        } else {
506                            scene["title"].set_string(title);
507                        }
508                    } else {
509                        let ori_title = scene["title"].as_str().unwrap_or("").to_string();
510                        while scene["title"].len() < self.language_index {
511                            scene["title"].push_member(ori_title.clone());
512                        }
513                        if self.insert_language {
514                            scene["title"].insert_member(self.language_index, title);
515                        } else {
516                            scene["title"].push_member(title);
517                        }
518                    }
519                    cur_mes = mes.next();
520                } else if scene["title"].is_list() {
521                    let m = match cur_mes {
522                        Some(m) => m,
523                        None => {
524                            return Err(anyhow::anyhow!(
525                                "No enough messages. (title at scene {i})"
526                            ));
527                        }
528                    };
529                    let mut title = m.message.clone();
530                    if let Some(replacement) = replacement {
531                        for (key, value) in replacement.map.iter() {
532                            title = title.replace(key, value);
533                        }
534                    }
535                    let ori_title = scene["title"][0].as_str().unwrap_or("").to_string();
536                    if self.insert_language {
537                        while scene["title"].len() < self.language_index {
538                            scene["title"].push_member(ori_title.clone());
539                        }
540                        scene["title"].insert_member(self.language_index, title);
541                    } else {
542                        while scene["title"].len() <= self.language_index {
543                            scene["title"].push_member(ori_title.clone());
544                        }
545                        scene["title"][self.language_index].set_string(title);
546                    }
547                    cur_mes = mes.next();
548                }
549            }
550            if scene["texts"].is_list() {
551                for (j, text) in scene["texts"].members_mut().enumerate() {
552                    if text.is_list() {
553                        if text.len() <= 1 {
554                            continue; // Skip if there are not enough values
555                        }
556                        if cur_mes.is_none() {
557                            cur_mes = mes.next();
558                        }
559                        if !text[0].is_string_or_null() {
560                            return Err(anyhow::anyhow!("name is not a string or null"));
561                        }
562                        let has_name = text[0].is_string();
563                        let has_display_name;
564                        if text[1].is_list() {
565                            if self.insert_language {
566                                let ori = text[1][0].clone();
567                                while text[1].len() < self.language_index {
568                                    text[1].push_member(ori.clone());
569                                }
570                                text[1].insert_member(self.language_index, ori.clone());
571                            } else {
572                                while text[1].len() <= self.language_index {
573                                    text[1][self.language_index] = text[1][0].clone();
574                                }
575                            }
576                            if text[1][self.language_index].is_list()
577                                && text[1][self.language_index].len() >= 2
578                            {
579                                if !text[1][self.language_index][0].is_string_or_null() {
580                                    return Err(anyhow::anyhow!(
581                                        "display name is not a string or null"
582                                    ));
583                                }
584                                if text[1][self.language_index][1].is_string() {
585                                    let m = match cur_mes.take() {
586                                        Some(m) => m,
587                                        None => {
588                                            return Err(anyhow::anyhow!(
589                                                "No enough messages. (text {j} at scene {i})"
590                                            ));
591                                        }
592                                    };
593                                    if has_name {
594                                        if let Some(name) = &m.name {
595                                            let mut name = name.clone();
596                                            if let Some(replacement) = replacement {
597                                                for (key, value) in replacement.map.iter() {
598                                                    name = name.replace(key, value);
599                                                }
600                                            }
601                                            text[1][self.language_index][0].set_string(name);
602                                        } else {
603                                            return Err(anyhow::anyhow!(
604                                                "Name is missing for message. (text {j} at scene {i}, message: {})",
605                                                m.message
606                                            ));
607                                        }
608                                    }
609                                    let mut message = m.message.clone();
610                                    if let Some(replacement) = replacement {
611                                        for (key, value) in replacement.map.iter() {
612                                            message = message.replace(key, value);
613                                        }
614                                    }
615                                    text[1][self.language_index][1]
616                                        .set_string(message.replace("\n", "\\n"));
617                                    // text length
618                                    text[1][self.language_index][2]
619                                        .set_i64(message.chars().count() as i64);
620                                    text[1][self.language_index][3]
621                                        .set_string(get_save_message(&message, true));
622                                    text[1][self.language_index][4]
623                                        .set_string(get_save_message(&message, false));
624                                }
625                            }
626                        } else {
627                            if text.len() <= 2 {
628                                continue; // Skip if there is no message
629                            }
630                            if !text[1].is_string_or_null() {
631                                return Err(anyhow::anyhow!(
632                                    "display name is not a string or null"
633                                ));
634                            }
635                            has_display_name = text[1].is_string();
636                            if text[2].is_string() {
637                                let m = match cur_mes.take() {
638                                    Some(m) => m,
639                                    None => {
640                                        return Err(anyhow::anyhow!(
641                                            "No enough messages.(text {j} at scene {i})"
642                                        ));
643                                    }
644                                };
645                                if has_name {
646                                    if let Some(name) = &m.name {
647                                        let mut name = name.clone();
648                                        if let Some(replacement) = replacement {
649                                            for (key, value) in replacement.map.iter() {
650                                                name = name.replace(key, value);
651                                            }
652                                        }
653                                        if has_display_name {
654                                            text[1].set_string(name);
655                                        } else {
656                                            text[0].set_string(name);
657                                        }
658                                    } else {
659                                        return Err(anyhow::anyhow!(
660                                            "Name is missing for message.(text {j} at scene {i})"
661                                        ));
662                                    }
663                                }
664                                let mut message = m.message.clone();
665                                if let Some(replacement) = replacement {
666                                    for (key, value) in replacement.map.iter() {
667                                        message = message.replace(key, value);
668                                    }
669                                }
670                                text[2].set_string(message.replace("\n", "\\n"));
671                            } else if text[2].is_list() {
672                                if self.insert_language {
673                                    let ori = text[2][0].clone();
674                                    while text[2].len() < self.language_index {
675                                        text[2].push_member(ori.clone());
676                                    }
677                                    text[2].insert_member(self.language_index, ori.clone());
678                                } else {
679                                    while text[2].len() <= self.language_index {
680                                        text[2][self.language_index] = text[2][0].clone();
681                                    }
682                                }
683                                if text[2][self.language_index].is_list()
684                                    && text[2][self.language_index].len() >= 2
685                                {
686                                    if !text[2][self.language_index][0].is_string_or_null() {
687                                        return Err(anyhow::anyhow!(
688                                            "display name is not a string or null"
689                                        ));
690                                    }
691                                    if text[2][self.language_index][1].is_string() {
692                                        let m = match cur_mes.take() {
693                                            Some(m) => m,
694                                            None => {
695                                                return Err(anyhow::anyhow!(
696                                                    "No enough messages.(text {j} at scene {i})"
697                                                ));
698                                            }
699                                        };
700                                        if has_name {
701                                            if let Some(name) = &m.name {
702                                                let mut name = name.clone();
703                                                if let Some(replacement) = replacement {
704                                                    for (key, value) in replacement.map.iter() {
705                                                        name = name.replace(key, value);
706                                                    }
707                                                }
708                                                text[2][self.language_index][0].set_string(name);
709                                            } else {
710                                                return Err(anyhow::anyhow!(
711                                                    "Name is missing for message.(text {j} at scene {i})"
712                                                ));
713                                            }
714                                        }
715                                        let mut message = m.message.clone();
716                                        if let Some(replacement) = replacement {
717                                            for (key, value) in replacement.map.iter() {
718                                                message = message.replace(key, value);
719                                            }
720                                        }
721                                        text[2][self.language_index][1]
722                                            .set_string(message.replace("\n", "\\n"));
723                                        text[2][self.language_index][2]
724                                            .set_i64(message.chars().count() as i64);
725                                        text[2][self.language_index][3]
726                                            .set_string(get_save_message(&message, true));
727                                        text[2][self.language_index][4]
728                                            .set_string(get_save_message(&message, false));
729                                    }
730                                }
731                            }
732                        }
733                    }
734                }
735            }
736            if scene["selects"].is_list() {
737                for select in scene["selects"].members_mut() {
738                    if select.is_object() {
739                        if cur_mes.is_none() {
740                            cur_mes = mes.next();
741                        }
742                        if self.language_index != 0
743                            && {
744                                if self.insert_language {
745                                    while select["language"].len() < self.language_index {
746                                        // TenShiSouZou
747                                        // first block is null
748                                        if select["language"].len() == 0 {
749                                            select["language"].push_member(PsbValueFixed::Null);
750                                            continue;
751                                        }
752                                        let mut obj = PsbObjectFixed::new();
753                                        obj["text"].set_str("");
754                                        obj["speechtext"].set_str("");
755                                        obj["searchtext"].set_str("");
756                                        obj["textlength"].set_i64(0);
757                                        select["language"][self.language_index].set_obj(obj);
758                                    }
759                                    let mut obj = PsbObjectFixed::new();
760                                    obj["text"].set_str("");
761                                    obj["speechtext"].set_str("");
762                                    obj["searchtext"].set_str("");
763                                    obj["textlength"].set_i64(0);
764                                    select["language"].insert_member(self.language_index, obj);
765                                } else {
766                                    while select["language"].len() <= self.language_index {
767                                        // TenShiSouZou
768                                        // first block is null
769                                        if select["language"].len() == 0 {
770                                            select["language"].push_member(PsbValueFixed::Null);
771                                            continue;
772                                        }
773                                        let mut obj = PsbObjectFixed::new();
774                                        obj["text"].set_str("");
775                                        obj["speechtext"].set_str("");
776                                        obj["searchtext"].set_str("");
777                                        obj["textlength"].set_i64(0);
778                                        select["language"][self.language_index].set_obj(obj);
779                                    }
780                                }
781                                true
782                            }
783                            && select["language"][self.language_index].is_object()
784                        {
785                            let lang_obj = &mut select["language"][self.language_index];
786                            if lang_obj["text"].is_string() {
787                                let m = match cur_mes.take() {
788                                    Some(m) => m,
789                                    None => {
790                                        return Err(anyhow::anyhow!("No enough messages."));
791                                    }
792                                };
793                                let mut text = m.message.clone();
794                                if let Some(replacement) = replacement {
795                                    for (key, value) in replacement.map.iter() {
796                                        text = text.replace(key, value);
797                                    }
798                                }
799                                lang_obj["text"].set_string(text.replace("\n", "\\n"));
800                                lang_obj["speechtext"].set_string(get_save_message(&text, true));
801                                lang_obj["searchtext"].set_string(get_save_message(&text, false));
802                                lang_obj["textlength"].set_i64(text.chars().count() as i64);
803                                continue;
804                            }
805                        } else if select["text"].is_string() {
806                            let m = match cur_mes.take() {
807                                Some(m) => m,
808                                None => {
809                                    return Err(anyhow::anyhow!("No enough messages."));
810                                }
811                            };
812                            let mut text = m.message.clone();
813                            if let Some(replacement) = replacement {
814                                for (key, value) in replacement.map.iter() {
815                                    text = text.replace(key, value);
816                                }
817                            }
818                            if self.insert_language {
819                                let ori_text = select["text"].as_str().unwrap_or("").to_string();
820                                let mut obj = PsbObjectFixed::new();
821                                obj["text"].set_string(ori_text.replace("\n", "\\n"));
822                                obj["speechtext"].set_string(get_save_message(&ori_text, true));
823                                obj["searchtext"].set_string(get_save_message(&ori_text, false));
824                                obj["textlength"].set_i64(ori_text.chars().count() as i64);
825                                if select["language"].len() < 1 {
826                                    select["language"].push_member(PsbValueFixed::Null);
827                                }
828                                select["language"].insert_member(1, obj);
829                            }
830                            select["text"].set_string(text.replace("\n", "\\n"));
831                        }
832                    }
833                }
834            }
835            comu.as_ref().map(|c| c.import(scene));
836        }
837        if cur_mes.is_some() || mes.next().is_some() {
838            return Err(anyhow::anyhow!("Some messages were not processed."));
839        }
840        let psb = psb.to_psb(true);
841        psb.finish_v4(file).map_err(|e| {
842            anyhow::anyhow!("Failed to write PSB to file {}: {:?}", self.filename, e)
843        })?;
844        Ok(())
845    }
846
847    fn custom_export(&self, filename: &Path, encoding: Encoding) -> Result<()> {
848        let s = if self.custom_yaml {
849            serde_yaml_ng::to_string(&self.psb)
850                .map_err(|e| anyhow::anyhow!("Failed to serialize to YAML: {}", e))?
851        } else {
852            json::stringify_pretty(self.psb.to_json(), 2)
853        };
854        let mut f = crate::utils::files::write_file(filename)?;
855        let b = encode_string(encoding, &s, false)?;
856        f.write_all(&b)?;
857        Ok(())
858    }
859
860    fn custom_import<'a>(
861        &'a self,
862        custom_filename: &'a str,
863        file: Box<dyn WriteSeek + 'a>,
864        _encoding: Encoding,
865        output_encoding: Encoding,
866    ) -> Result<()> {
867        let data = crate::utils::files::read_file(custom_filename)?;
868        let s = decode_to_string(output_encoding, &data, true)?;
869        let psb = if self.custom_yaml {
870            let data: VirtualPsbFixedData = serde_yaml_ng::from_str(&s)
871                .map_err(|e| anyhow::anyhow!("Failed to deserialize YAML: {}", e))?;
872            let mut psb = self.psb.clone();
873            psb.set_data(data);
874            psb.to_psb(true)
875        } else {
876            let json = json::parse(&s)?;
877            let mut psb = self.psb.clone();
878            psb.from_json(&json)?;
879            psb.to_psb(true)
880        };
881        psb.finish_v4(file).map_err(|e| {
882            anyhow::anyhow!("Failed to write PSB to file {}: {:?}", self.filename, e)
883        })?;
884        Ok(())
885    }
886}
887
888#[derive(Debug)]
889struct ExportMes {
890    pub messages: HashSet<String>,
891    pub key: HashSet<String>,
892    text_key: String,
893}
894
895impl ExportMes {
896    pub fn new(key: Vec<String>, language: Option<String>) -> Self {
897        Self {
898            messages: HashSet::new(),
899            key: HashSet::from_iter(key.into_iter()),
900            text_key: language.map_or_else(|| String::from("text"), |s| format!("text_{}", s)),
901        }
902    }
903
904    pub fn export(&mut self, value: &PsbValueFixed) {
905        match value {
906            PsbValueFixed::Object(obj) => {
907                for (k, v) in obj.iter() {
908                    if self.key.contains(k) {
909                        if let PsbValueFixed::List(list) = v {
910                            for item in list.iter() {
911                                if let PsbValueFixed::Object(obj) = item {
912                                    if let Some(s) = obj[&self.text_key].as_str() {
913                                        self.messages.insert(s.replace("\\n", "\n"));
914                                    } else if let Some(s) = obj["text"].as_str() {
915                                        self.messages.insert(s.replace("\\n", "\n"));
916                                    }
917                                }
918                            }
919                        }
920                    } else {
921                        self.export(v);
922                    }
923                }
924            }
925            PsbValueFixed::List(list) => {
926                let list = list.values();
927                if list.len() > 1 {
928                    if let PsbValueFixed::String(s) = &list[0] {
929                        if self.key.contains(s.string()) {
930                            for i in 1..list.len() {
931                                if let PsbValueFixed::String(s) = &list[i - 1] {
932                                    if s.string() == &self.text_key {
933                                        if let PsbValueFixed::String(text) = &list[i] {
934                                            self.messages
935                                                .insert(text.string().replace("\\n", "\n"));
936                                        }
937                                    }
938                                }
939                            }
940                            if self.text_key == "text" {
941                                return;
942                            }
943                            for i in 1..list.len() {
944                                if let PsbValueFixed::String(s) = &list[i - 1] {
945                                    if s.string() == "text" {
946                                        if let PsbValueFixed::String(text) = &list[i] {
947                                            self.messages
948                                                .insert(text.string().replace("\\n", "\n"));
949                                        }
950                                    }
951                                }
952                            }
953                            return;
954                        }
955                    }
956                }
957                for item in list {
958                    self.export(item);
959                }
960            }
961            _ => {}
962        }
963    }
964}
965
966lazy_static::lazy_static! {
967    static ref DUP_WARN_SHOWN: Mutex<HashSet<(String, usize, String)>> = Mutex::new(HashSet::new());
968    static ref NOT_FOUND_WARN_SHOWN: Mutex<HashSet<String>> = Mutex::new(HashSet::new());
969}
970
971fn warn_dup(original: String, count: usize, filename: String) {
972    let mut guard = DUP_WARN_SHOWN.lock_blocking();
973    if guard.contains(&(original.clone(), count, filename.clone())) {
974        return;
975    }
976    eprintln!(
977        "Warning: chat message '{}' has {} duplicates in translation table '{}'. Using the first one.",
978        original, count, filename
979    );
980    crate::COUNTER.inc_warning();
981    guard.insert((original.clone(), count, filename.clone()));
982}
983
984fn warn_not_found(original: String) {
985    let mut guard = NOT_FOUND_WARN_SHOWN.lock_blocking();
986    if guard.contains(&original) {
987        return;
988    }
989    eprintln!(
990        "Warning: chat message '{}' not found in translation table.",
991        original
992    );
993    crate::COUNTER.inc_warning();
994    guard.insert(original);
995}
996
997#[derive(Debug)]
998struct ImportMes<'a> {
999    messages: &'a Arc<HashMap<String, HashMap<String, (String, usize)>>>,
1000    replacement: Option<&'a ReplacementTable>,
1001    key: HashSet<String>,
1002    text_key: String,
1003    filename: String,
1004    ori_text_key: Option<String>,
1005}
1006
1007impl<'a> ImportMes<'a> {
1008    pub fn new(
1009        messages: &'a Arc<HashMap<String, HashMap<String, (String, usize)>>>,
1010        replacement: Option<&'a ReplacementTable>,
1011        key: Vec<String>,
1012        lang: Option<String>,
1013        filename: String,
1014        ori_lang: Option<String>,
1015    ) -> Self {
1016        Self {
1017            messages,
1018            replacement,
1019            key: HashSet::from_iter(key.into_iter()),
1020            text_key: lang.map_or_else(|| String::from("text"), |s| format!("text_{}", s)),
1021            filename: std::path::Path::new(&filename)
1022                .file_stem()
1023                .map(|s| s.to_string_lossy().to_string())
1024                .unwrap_or_else(|| "global".to_string()),
1025            ori_text_key: ori_lang.map(|s| format!("text_{}", s)),
1026        }
1027    }
1028
1029    fn get_message(&self, original: &str) -> Option<String> {
1030        if let Some(global) = self.messages.get(&self.filename) {
1031            if let Some(text) = global.get(original) {
1032                if text.1 > 1 {
1033                    warn_dup(original.to_string(), text.1, self.filename.clone());
1034                }
1035                return Some(text.0.clone());
1036            }
1037        }
1038        if self.filename == "global" {
1039            return None;
1040        }
1041        if let Some(file) = self.messages.get("global") {
1042            if let Some(text) = file.get(original) {
1043                if text.1 > 1 {
1044                    warn_dup(original.to_string(), text.1, "global".to_string());
1045                }
1046                return Some(text.0.clone());
1047            }
1048        }
1049        None
1050    }
1051
1052    pub fn import(&self, value: &mut PsbValueFixed) {
1053        match value {
1054            PsbValueFixed::Object(obj) => {
1055                for (k, v) in obj.iter_mut() {
1056                    if self.key.contains(k) {
1057                        for obj in v.members_mut() {
1058                            if let Some(text) = obj[&self.text_key].as_str() {
1059                                if let Some(replace_text) = self.get_message(text) {
1060                                    let mut text = replace_text.clone();
1061                                    if let Some(replacement) = self.replacement {
1062                                        for (key, value) in replacement.map.iter() {
1063                                            text = text.replace(key, value);
1064                                        }
1065                                    }
1066                                    if self.text_key == "text" {
1067                                        if let Some(ori_key) = &self.ori_text_key {
1068                                            let ori_text =
1069                                                obj["text"].as_str().unwrap_or("").to_string();
1070                                            obj[ori_key].set_string(ori_text);
1071                                        }
1072                                    }
1073                                    obj[&self.text_key].set_string(text.replace("\n", "\\n"));
1074                                    continue;
1075                                } else {
1076                                    warn_not_found(text.to_string());
1077                                }
1078                            }
1079                            if let Some(text) = obj["text"].as_str() {
1080                                if let Some(replace_text) = self.get_message(text) {
1081                                    let mut text = replace_text.clone();
1082                                    if let Some(replacement) = self.replacement {
1083                                        for (key, value) in replacement.map.iter() {
1084                                            text = text.replace(key, value);
1085                                        }
1086                                    }
1087                                    if let Some(ori_key) = &self.ori_text_key {
1088                                        let ori_text =
1089                                            obj["text"].as_str().unwrap_or("").to_string();
1090                                        obj[ori_key].set_string(ori_text);
1091                                    }
1092                                    obj[&self.text_key].set_string(text.replace("\n", "\\n"));
1093                                } else {
1094                                    warn_not_found(text.to_string());
1095                                }
1096                            }
1097                        }
1098                    } else {
1099                        self.import(v);
1100                    }
1101                }
1102            }
1103            PsbValueFixed::List(list) => {
1104                if list.len() > 1 {
1105                    if list[0].as_str().map_or(false, |s| self.key.contains(s)) {
1106                        for i in 1..list.len() {
1107                            if list[i - 1] == self.text_key {
1108                                if let Some(text) = list[i].as_str() {
1109                                    if let Some(replace_text) = self.get_message(text) {
1110                                        let mut text = replace_text.clone();
1111                                        if let Some(replacement) = self.replacement {
1112                                            for (key, value) in replacement.map.iter() {
1113                                                text = text.replace(key, value);
1114                                            }
1115                                        }
1116                                        if self.text_key == "text" {
1117                                            if let Some(ori_key) = &self.ori_text_key {
1118                                                let len = list.len();
1119                                                let ori_text =
1120                                                    list[i].as_str().unwrap_or("").to_string();
1121                                                list[len].set_str(ori_key);
1122                                                list[len + 1].set_string(ori_text);
1123                                            }
1124                                        }
1125                                        list[i].set_string(text.replace("\n", "\\n"));
1126                                        return;
1127                                    } else {
1128                                        warn_not_found(text.to_string());
1129                                    }
1130                                }
1131                            }
1132                        }
1133                        if self.text_key == "text" {
1134                            return;
1135                        }
1136                        for i in 1..list.len() {
1137                            if list[i - 1] == "text" {
1138                                if let Some(text) = list[i].as_str() {
1139                                    if let Some(replace_text) = self.get_message(text) {
1140                                        let mut text = replace_text.clone();
1141                                        if let Some(replacement) = self.replacement {
1142                                            for (key, value) in replacement.map.iter() {
1143                                                text = text.replace(key, value);
1144                                            }
1145                                        }
1146                                        let len = list.len();
1147                                        list[len].set_str(&self.text_key);
1148                                        list[len + 1].set_string(text.replace("\n", "\\n"));
1149                                        return;
1150                                    } else {
1151                                        warn_not_found(text.to_string());
1152                                    }
1153                                }
1154                            }
1155                        }
1156                        return;
1157                    }
1158                }
1159                for item in list.iter_mut() {
1160                    self.import(item);
1161                }
1162            }
1163            _ => {}
1164        }
1165    }
1166}
1167
1168lazy_static::lazy_static! {
1169    static ref CONTROL: Regex = Regex::new("%[^%;]*;").unwrap();
1170    static ref RUBY: Regex = Regex::new(r"(?<!\\)\[([^\]]*)\](.?)").unwrap();
1171    static ref COLOR: Regex = Regex::new(r"#[0-9a-fA-F]{6,8};").unwrap();
1172}
1173
1174fn get_save_message(s: &str, in_ruby: bool) -> String {
1175    let mut s = s.replace("\n", "");
1176    s = CONTROL.replace_all(&s, "").to_string();
1177    s = COLOR.replace_all(&s, "").to_string();
1178    s = RUBY
1179        .replace_all(&s, if in_ruby { "$1" } else { "$2" })
1180        .to_string();
1181    s.replace("%r", "").replace("\\[", "[")
1182}
1183
1184#[test]
1185fn test_get_save_message() {
1186    let s = "%n;Test\n[ruby]测[test\\]试%ok;[ok]";
1187    assert_eq!(get_save_message(s, true), "Testrubytest\\ok");
1188    assert_eq!(get_save_message(s, false), "Test测试");
1189    let another = "[Start]a";
1190    assert_eq!(get_save_message(another, true), "Start");
1191    assert_eq!(get_save_message(another, false), "a");
1192    let escaped = "\\[Start]a";
1193    assert_eq!(get_save_message(escaped, true), "[Start]a");
1194    assert_eq!(get_save_message(escaped, false), "[Start]a");
1195    let real_word = "「こんな、感じとか、ですか……? うっふ~ん……%f$ハート$;#00ffadd6;♥%r」";
1196    assert_eq!(
1197        get_save_message(real_word, true),
1198        "「こんな、感じとか、ですか……? うっふ~ん……♥」"
1199    );
1200    let s = "「あっは%f$ハート$;#00ffadd6;♥%r 凄い出してくれてる%f$ハート$;#00ffadd6;♥%r」";
1201    assert_eq!(
1202        get_save_message(s, true),
1203        "「あっは♥ 凄い出してくれてる♥」"
1204    );
1205}